import torch
from torch import nn
import numpy


class FM(nn.Module):
    def __init__(self, n=10, k=2):
        super(FM, self).__init__()
        self.n = n
        self.k = k
        self.linear = nn.Linear(self.n, 1)   # 前两项线性层
        self.V = nn.Parameter(torch.randn(self.n, self.k))   # 交互矩阵
        # self.relu = nn.ReLU()

    def fm_layer(self, x):
        linear_part = self.linear(x)
        interaction_part_1 = torch.mm(x, self.V)
        interaction_part_1 = torch.pow(interaction_part_1, 2)
        interaction_part_2 = torch.mm(torch.pow(x, 2), torch.pow(self.V, 2))

        fm_interaction = interaction_part_1 - interaction_part_2
        fm_interaction = torch.sum(fm_interaction, dim=1, keepdim=True)
        fm_interaction = 0.5 * fm_interaction
        # output = linear_part + torch.sum(0.5 * interaction_part_2 - interaction_part_1)
        output = linear_part + fm_interaction
        # 添加映射
        # output = self.relu(output)
        return output

    def forward(self, x):
        return self.fm_layer(x)


fm = FM(10, 5)
x = torch.randn(2,10)
output = fm(x)
print(output)